import pandas as pd
import pymysql
import os
from sqlalchemy import create_engine

# Database connection parameters - replace with your actual credentials
DB_HOST = ''
DB_USER = 'zazhi4'
DB_PASSWORD = ''
DB_NAME = 'zazhi4'
DB_PORT = 3307


# Excel file path - replace with your actual file path
EXCEL_FILE_PATH = './ejia提前一个月的.xlsx'

def update_database():
    try:
        # Create database connection
        connection = pymysql.connect(
            host=DB_HOST,
            user=DB_USER,
            password=DB_PASSWORD,
            database=DB_NAME,
            charset='utf8mb4',
            port=DB_PORT,
            cursorclass=pymysql.cursors.DictCursor
        )
        
        print("Connected to database successfully!")
        
        try:
            with connection.cursor() as cursor:
                # First update: Set all records to is_two_or_three = 2
                reset_sql = "UPDATE fa_shop_goods SET is_two_or_three = 2"
                cursor.execute(reset_sql)
                print(f"Reset all records to is_two_or_three = 2")
                
                # Read Excel file - assuming first column contains goods_name values
                if not os.path.exists(EXCEL_FILE_PATH):
                    print(f"Error: Excel file not found at {EXCEL_FILE_PATH}")
                    return
                
                df = pd.read_excel(EXCEL_FILE_PATH)
                
                # Get the name of the first column
                first_column = df.columns[0]
                
                # Extract goods_name values from first column
                goods_names = df[first_column].dropna().tolist()
                print(f"Read {len(goods_names)} goods_name values from Excel file")
                
                # Update in batches to avoid query too long errors
                batch_size = 50
                total_updated = 0
                
                for i in range(0, len(goods_names), batch_size):
                    batch = goods_names[i:i+batch_size]
                    
                    # Create placeholders for SQL query
                    placeholders = ', '.join(['%s'] * len(batch))
                    
                    # Update matching records to is_two_or_three = 1
                    update_sql = f"""
                    UPDATE fa_shop_goods 
                    SET is_two_or_three = 1 
                    WHERE title IN ({placeholders})
                    """
                    
                    cursor.execute(update_sql, batch)
                    affected_rows = cursor.rowcount
                    total_updated += affected_rows
                    print(f"Batch {i//batch_size + 1}: Updated {affected_rows} records to is_two_or_three = 1")
                
                # Commit the transaction
                connection.commit()
                
                # Get final counts
                cursor.execute("SELECT COUNT(*) as count FROM fa_shop_goods WHERE is_two_or_three = 1")
                count_result = cursor.fetchone()
                print(f"Final count of records with is_two_or_three = 1: {count_result['count']}")
                
                print("Database update completed successfully!")
                
        finally:
            connection.close()
            print("Database connection closed")
            
    except Exception as e:
        print(f"Error: {str(e)}")

if __name__ == "__main__":
    update_database()

